from Crypto.Cipher import AES

class AESCMAC:
    """AES-CMAC算法工具类 (RFC4493标准实现)"""
    
    @staticmethod
    def xor_bytes(a, b):
        """字节数组异或操作"""
        return bytes(x ^ y for x, y in zip(a, b))
    
    @staticmethod
    def left_shift(block):
        """左移1位操作"""
        output = bytearray(block)
        carry = 0
        for i in reversed(range(len(output))):
            new_carry = (output[i] & 0x80) >> 7
            output[i] = ((output[i] << 1) | carry) & 0xFF
            carry = new_carry
        return bytes(output)
    
    @staticmethod
    def generate_subkeys(key):
        """生成CMAC子密钥K1和K2"""
        cipher = AES.new(key, AES.MODE_ECB)
        zero_block = bytes(16)
        L = cipher.encrypt(zero_block)
        
        # 生成K1
        K1 = AESCMAC.left_shift(L)
        if L[0] & 0x80:
            K1 = AESCMAC.xor_bytes(K1, bytes.fromhex('00000000000000000000000000000087'))
            
        # 生成K2
        K2 = AESCMAC.left_shift(K1)
        if K1[0] & 0x80:
            K2 = AESCMAC.xor_bytes(K2, bytes.fromhex('00000000000000000000000000000087'))
            
        return K1, K2

    @staticmethod
    def calculate(data, key = '2b7e1516 28aed2a6 abf71588 09cf4f3c'):
        """
        计算AES-CMAC值
        :param data: 输入数据(hex字符串或bytes)
        :param key: 16/24/32字节密钥(hex字符串或bytes)
        :return: CMAC结果(16字节bytes)
        """
        # 处理密钥输入
        if isinstance(key, str):
            key_bytes = bytes.fromhex(key.replace(" ", ""))
        else:
            key_bytes = bytes(key)
            
        # 验证密钥长度
        if len(key_bytes) not in (16, 24, 32):
            raise ValueError("密钥长度必须为16/24/32字节")
            
        # 处理数据输入
        if isinstance(data, str):
            data_bytes = bytes.fromhex(data.replace(" ", ""))
        else:
            data_bytes = bytes(data)
        
        # RFC4493标准填充
        pad_length = 16 - (len(data_bytes) % 16) if len(data_bytes) % 16 != 0 else 0
        if pad_length > 0:
            data_bytes += b'\x80' + b'\x00'*(pad_length-1)
        
        K1, K2 = AESCMAC.generate_subkeys(key_bytes)
        cipher = AES.new(key_bytes, AES.MODE_CBC, iv=bytes(16))
        
        # 修改数据块处理逻辑
        blocks = [data_bytes[i:i+16] for i in range(0, len(data_bytes), 16)]
        if len(blocks) == 0:
            blocks = [bytes(16)]
        
        last_block = blocks[-1]
        cipher = AES.new(key_bytes, AES.MODE_CBC, iv=bytes(16))
        
        # 处理完整块
        for block in blocks[:-1]:
            cipher.encrypt(block)
        
        # 处理最后一个块
        if len(last_block) < 16:
            last_block = AESCMAC.xor_bytes(last_block + b'\x80' + bytes(16 - len(last_block) - 1), K1)
        else:
            last_block = AESCMAC.xor_bytes(last_block, K2)
        
        return cipher.encrypt(last_block)

    @staticmethod 
    def calculate_hex(data, key = '2b7e1516 28aed2a6 abf71588 09cf4f3c'):
        """
        计算AES-CMAC值(hex字符串形式)
        :param data: 输入数据(hex字符串或bytes)
        :param key: 16/24/32字节密钥(hex字符串或bytes)
        :return: CMAC结果(32字符hex字符串)
        """
        return AESCMAC.calculate(data, key).hex().upper()

    @staticmethod
    def verify(data, key = '2b7e1516 28aed2a6 abf71588 09cf4f3c', expected_cmac = ''):
        """
        验证AES-CMAC值
        :param data: 输入数据(hex字符串或bytes)
        :param key: 16/24/32字节密钥(hex字符串或bytes)
        :param expected_cmac: 预期CMAC值(hex字符串或bytes)
        :return: bool 验证结果
        """
        actual_cmac = AESCMAC.calculate(data, key)
        if isinstance(expected_cmac, str):
            expected_bytes = bytes.fromhex(expected_cmac.replace(" ", ""))
        else:
            expected_bytes = bytes(expected_cmac)
        return actual_cmac == expected_bytes


# 测试代码
if __name__ == "__main__":
    # 测试用例
    test_vectors = [
        {
            "key": "ae 63 45 de 6a ba d0 ec a8 f1 a5 26 22 0d 5e 08",
            "data": "65 be b3 cf 61 8c db 5d",
            "cmac": "25858fc41b1a9bf735d70e6f59daaabc"
        },
        {
            "key": "d5 4c 3f 8a 9b 1e 20 7d 6a 48 72 c3 9f 0e 5b 1a",
            "data": "31 32 33 34 35 36 37 38 39 30 31 32 33 34 35 36 37",
            "cmac": "ae 63 45 de 6a ba d0 ec a8 f1 a5 26 22 0d 5e 08"
        },
        # {
        #     "key": "01020304 05060708 09000102 03040506",
        #     "data": "6bc1bee2 2e409f96 e93d7e11 7393172a",
        #     "cmac": "E18AD597 32AE1029 23DF6401 3079C413"
        # },
        # {
        #     "key": "01020304 05060708 09000102 03040506",
        #     "data": "6bc1bee2 2e409f96 e93d7e11 7393172a ae2d8a57 1e03ac9c 9eb76fac 45af8e51 30c81c46 a35ce411",
        #     "cmac": "1861D07B A5E96229 EED61110 B45E5845"
        # }

    ]

    for i, vector in enumerate(test_vectors, 1):
        print(f"\n测试用例 {i}:")
        cmac = AESCMAC.calculate_hex(vector["data"], vector["key"])
        print(f"密钥: {vector['key'].replace(' ', '')}")
        print(f"数据: {vector['data'].replace(' ', '')}")
        print(f"计算结果: {cmac}")
        print(f"预期结果: {vector['cmac'].replace(' ', '')}")
        print(f"验证结果: {'通过' if AESCMAC.verify(vector['data'], vector['key'], vector['cmac']) else '失败'}")

